#!/usr/bin/env python3
import numpy as np
import sys
import os 
import matplotlib.pyplot as plt
import scienceplots
import glob
import json
from scipy.spatial.transform import Rotation as R

from my_common_lib import linestyle, marker, labelfont, titlefont, cm, color





def load_json_file(file_path):
    """
    Load data from a JSON file.
    """
    with open(file_path, 'r') as f:
        return json.load(f)

def plot_odom_data(data, save_dir):
    """
    Plot position and velocity from Odometry data.
    """
    time = [entry['time'] for entry in data]
    
    # Plot position
    with plt.style.context(['science','ieee']):
        plt.figure(1, figsize=(8.5 * cm, 6 * cm), dpi=600)
        plt.plot(time, [entry['p_x'] for entry in data], label='Position X', color='red')
        plt.plot(time, [entry['p_y'] for entry in data], label='Position Y', color='green')
        plt.plot(time, [entry['p_z'] for entry in data], label='Position Z', color='blue')
        plt.title('Odometry Position Over Time')
        plt.xlabel('Time (s)')
        plt.ylabel('Position (m)')
        plt.legend()
        plt.grid()
        plt.savefig(os.path.join(save_dir, 'odometry_position.png'), dpi=600)
        plt.close()
    
        # Plot linear velocity
        plt.figure(2, figsize=(8.5 * cm, 6 * cm), dpi=600)
        plt.plot(time, [entry['vel_linear_x'] for entry in data], label='Linear Velocity X', color='red')
        plt.plot(time, [entry['vel_linear_y'] for entry in data], label='Linear Velocity Y', color='green')
        plt.plot(time, [entry['vel_linear_z'] for entry in data], label='Linear Velocity Z', color='blue')
        plt.title('Odometry Linear Velocity Over Time')
        plt.xlabel('Time (s)')
        plt.ylabel('Linear Velocity (m/s)')
        plt.legend()
        plt.grid()
        plt.savefig(os.path.join(save_dir, 'odometry_velocity.png'), dpi=600)
        plt.close()

def plot_nodeframe_data(data, save_dir):
    """
    Plot distance and signal strength for each node in LinktrackNodeframe3 data.
    """
    for i, entry in enumerate(data):
        nodes = entry['nodes']
        distances = [node['distance'] for node in nodes]
        fp_rssi = [node['fp_rssi'] for node in nodes]
        with plt.style.context(['science','ieee']):
            plt.figure(3, figsize=(8.5 * cm, 6 * cm), dpi=600)
            plt.scatter(distances, fp_rssi, label='Timestamp {}'.format(entry["local_time"]))
            plt.title('Node Distance vs Signal Strength (Entry {})'.format(i))
            plt.xlabel('Distance (m)')
            plt.ylabel('FP RSSI (dBm)')
            plt.legend()
            plt.grid()
            plt.savefig(os.path.join(save_dir, 'nodeframe_distance_rssi_{}.png'.format(i)), dpi=600)
            plt.close()

def plot_imu_data(data, save_dir):
    """
    Plot orientation from IMU data.
    """
    time = [entry['time'] for entry in data]
    with plt.style.context(['science','ieee']):
        plt.figure(4, figsize=(8.5 * cm, 6 * cm), dpi=600)
        plt.plot(time, [entry['orient_x'] for entry in data], label='Orientation X', color='red')
        plt.plot(time, [entry['orient_y'] for entry in data], label='Orientation Y', color='green')
        plt.plot(time, [entry['orient_z'] for entry in data], label='Orientation Z', color='blue')
        plt.plot(time, [entry['orient_w'] for entry in data], label='Orientation W', color='purple')
        plt.title('IMU Orientation Over Time')
        plt.xlabel('Time (s)')
        plt.ylabel('Orientation (quaternion)')
        plt.legend()
        plt.grid()
        plt.savefig(os.path.join(save_dir, 'imu_orientation.png'), dpi=600)
        plt.close()

def plot_one_dimension_data(vec_x, vec_y, vec_label, title, xlabel, ylabel, save_folder, file_name, show = False, colors = None):

    with plt.style.context(['science','ieee']):
        plt.figure(4, figsize=(8.5 * cm, 6 * cm), dpi=600)
        if colors == None:
            colors = color
        for i in range(len(vec_x)):
            plt.plot(vec_x[i], vec_y[i], label=vec_label[i], color = colors[i])
        # plt.title(title)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.legend()
        plt.grid()
        if show == True:
            plt.show()
        else:
            plt.savefig(os.path.join(save_folder, file_name), dpi=600)
            plt.close()
            print("Plot saved in {}".format(os.path.join(save_folder, file_name)))

def plot_one_dimension_data_two_pair(vec1, vec2, title, xlabel, ylabel, save_folder, file_name, show = False, colors = None):

    vec1_x, vec1_y, vec1_label = vec1[0], vec1[1], vec1[2]
    vec2_x, vec2_y, vec2_label = vec2[0], vec2[1], vec2[2]
    with plt.style.context(['science','ieee']):
        plt.figure(4, figsize=(8.5 * cm, 6 * cm), dpi=600)
        if colors == None:
            colors = color
        for i in range(len(vec1_x)):
            plt.plot(vec1_x[i], vec1_y[i], label=vec1_label[i], color = colors[i], linestyle = linestyle[0])
        for i in range(len(vec2_x)):
            plt.plot(vec2_x[i], vec2_y[i], label=vec2_label[i], color = colors[i], linestyle = linestyle[1])
        # plt.title(title)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.legend()
        plt.grid()
        if show == True:
            plt.show()
        else:
            plt.savefig(os.path.join(save_folder, file_name), dpi=600)
            plt.close()
            print("Plot saved in {}".format(os.path.join(save_folder, file_name)))

def process_json_files(json_path, topics, types):
    """
    Process JSON files in a directory and generate plots.
    """
    if not os.path.exists(json_path):
        raise FileExistsError

    for file_name in os.listdir(json_path):
        file_path = os.path.join(json_path, file_name)
        if not file_name.endswith('.json'):
            continue

        data = load_json_file(file_path)
        print("Processing file: {}".format(file_name))

        if 'odom' in file_name:
            plot_odom_data(data, json_path)
        elif 'nodeframe3' in file_name:
            plot_nodeframe_data(data, json_path)
        elif 'imu' in file_name:
            plot_imu_data(data, json_path)

def load_orientation_data(file_path):
    data = load_json_file(file_path)
    array_data = [
        [entry["time"], entry["orient_x"], entry["orient_y"], entry["orient_z"], entry["orient_w"]] for entry in data
    ]
    # Convert the list of lists into a NumPy array
    return np.array(array_data)

def load_position_data(file_path):
    data = load_json_file(file_path)
    array_data = [
        [entry["time"], entry["p_x"], entry["p_y"], entry["p_z"]] for entry in data
    ]
    # Convert the list of lists into a NumPy array
    return np.array(array_data)

def load_nodeframe3_data(file_path):
    data = load_json_file(file_path)
    array_data = [[],[],[],[],[],[]]
    for entry in data:
        for node in entry["nodes"]:
            try:
                if node["distance"] < 20:
                    array_data[node["id"]].append([entry["time"], node["distance"]])
                else:
                    raise Exception("distance too large: ", node["distance"])
            except Exception as e:
                print("arise error", e, "node_id: ", node["id"])
    # Convert the list of lists into a NumPy array
    return array_data



   
       
        